package com.intmall.flink.process;

import com.intmall.flink.operator.ClickSource;
import com.intmall.flink.operator.Event;
import com.intmall.flink.window.UrlViewCount;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import java.sql.Timestamp;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Comparator;

public class TopNExample {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        SingleOutputStreamOperator<Event> stream = env.addSource(new ClickSource())
                .assignTimestampsAndWatermarks(WatermarkStrategy.<Event>forBoundedOutOfOrderness(Duration.ZERO)
                        .withTimestampAssigner(new SerializableTimestampAssigner<Event>() {
                            @Override
                            public long extractTimestamp(Event element, long recordTimestamp) {
                                return element.getTimestamp();
                            }
                        }));

        // 1.按照url分组，统计窗口内每个url的统计量
        SingleOutputStreamOperator<UrlViewCount> urlCountStream = stream.keyBy(data -> data.getUrl())
                .window(TumblingEventTimeWindows.of(Time.seconds(10), Time.seconds(5)))
                .aggregate(new AggregateFunction<Event, Long, Long>() {

                               @Override
                               public Long createAccumulator() {
                                   return 0L;
                               }

                               @Override
                               public Long add(Event value, Long accumulator) {
                                   return accumulator + 1;
                               }

                               @Override
                               public Long getResult(Long accumulator) {
                                   return accumulator;
                               }

                               @Override
                               public Long merge(Long a, Long b) {
                                   return null;
                               }
                           }, new ProcessWindowFunction<Long, UrlViewCount, String, TimeWindow>() {
                               @Override
                               public void process(String s, ProcessWindowFunction<Long, UrlViewCount, String, TimeWindow>.Context context, Iterable<Long> elements, Collector<UrlViewCount> out) throws Exception {
                                   long start = context.window().getStart();
                                   long end = context.window().getEnd();
                                   out.collect(new UrlViewCount(s, elements.iterator().next(), start, end));
                               }
                           }
                );

        urlCountStream.print("url count");

        // 2. 对于同一窗口统计出的访问量，进行收集和排序
        urlCountStream.keyBy(data -> data.getWindowEnd())
                .process(new TopNProcessResult(2));
        env.execute();
    }

    public static class TopNProcessResult extends KeyedProcessFunction<Long, UrlViewCount, String> {
        private Integer num;
        // 定义列表状态
        private ListState<UrlViewCount> urlViewCountListState;

        public TopNProcessResult(Integer num) {
            this.num = num;
        }

        @Override
        public void open(Configuration parameters) throws Exception {
            urlViewCountListState = getRuntimeContext().getListState(new ListStateDescriptor<UrlViewCount>("url-count-list",
                    Types.POJO(UrlViewCount.class)));
        }

        @Override
        public void processElement(UrlViewCount value, KeyedProcessFunction<Long, UrlViewCount, String>.Context ctx, Collector<String> out) throws Exception {
            urlViewCountListState.add(value);
            // 定时器 WindowEnd + 1毫秒
            ctx.timerService().registerEventTimeTimer(ctx.getCurrentKey() + 1);
        }

        @Override
        public void onTimer(long timestamp, KeyedProcessFunction<Long, UrlViewCount, String>.OnTimerContext ctx, Collector<String> out) throws Exception {
            // 定时器到期后的操作
            ArrayList<UrlViewCount> urlViewCounts = new ArrayList<>();
            for (UrlViewCount urlViewCount : urlViewCountListState.get()) {
                urlViewCounts.add(urlViewCount);
            }

            urlViewCounts.sort(new Comparator<UrlViewCount>() {
                @Override
                public int compare(UrlViewCount o1, UrlViewCount o2) {
                    return (int) (o2.getCount() - o1.getCount());
                }
            });

            StringBuilder result = new StringBuilder();
            result.append("窗口结束时间：" + new Timestamp(ctx.getCurrentKey())).append("\n");
            for (int i = 0; i < 2; i++) {
                UrlViewCount urlViewCount = urlViewCounts.get(i);
                result.append("No.").append(i + 1).append(" url: ")
                        .append(urlViewCount.getUrl()).append(" 访问量：").append(urlViewCount.getCount()).append("\n");
            }
            System.out.println(result);
        }
    }
}
